package beast.evolution.operators;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

import org.apache.commons.math.random.RandomDataImpl;

import beast.core.Description;
import beast.core.Input;
import beast.core.Operator;
import beast.evolution.tree.Node;
import beast.evolution.tree.Tree3;

@Description("Scales a parameter or a complete beast.tree (depending on which of the two is specified.")
public class PathSampleOperator extends Operator {

	RandomDataImpl randGenerator = new RandomDataImpl();
	Random generator = new Random();

	public Input<Tree3> m_pTree = new Input<Tree3>("tree3",
			"for now, we use entire tree as input");

	int sequenceSize;

	@Override
	public void initAndValidate() {
		sequenceSize = m_pTree.get().m_alignment.get().getSiteCount();
	}

	@Override
	public double proposal() {
		// hard code to test utility for now
		NielsenSampler(m_pTree.get().getNode(2));
		return 0;
	}

	public void NielsenSampler(Node currentNode) {
		// System.out.println("current species is:" + currentNode.getID()
		// + " and " + "current branch is:" + currentNode.getLength());

		// total number for testing purpose
		int totalNum = 0;

		// loop over all sites
		for (int siteNum = 0; siteNum < sequenceSize; siteNum++) {
			totalNum = totalNum + NielsenOneSiteSampler(currentNode, siteNum);
		}

		System.out.println("average number of substitution is:"
				+ (double) totalNum / (double) sequenceSize);
	}

	public int NielsenOneSiteSampler(Node currentNode, int siteNum) {
		// System.out.println("current species is:" + currentNode.getID()
		// + " and " + "current branch is:" + currentNode.getLength()
		// + " and siteNum is: " + siteNum);

		int childState = m_pTree.get().getSequenceByNr(currentNode.getNr())[siteNum];
		int parentState = m_pTree.get().getSequenceByNr(
				currentNode.getParent().getNr())[siteNum];

		// System.out.println("the childState is: " + childState
		// + " the parentState is: " + parentState);

		// initialize all variables
		double totalTime = currentNode.getLength();
		int endState = -1;
		double currentTime = 0;
		int currentState = parentState;

		// store sampled path and time
		List<Integer> sampledPath = new ArrayList<Integer>();
		List<Double> sampledTime = new ArrayList<Double>();

		// timeInterval that needs to be sampled in each round
		double timeInterval = 0;

		// the work horse
		if (parentState == childState) {
			while (endState != childState) {
				// initialize path
				sampledPath.clear();
				sampledPath.add(new Integer(parentState));
				sampledTime.clear();
				sampledTime.add(new Double(0.0));

				// initialize time and state
				currentTime = 0.0;
				currentState = parentState;

				while (currentTime < totalTime) {
					timeInterval = randGenerator.nextExponential(1.0);
					// System.out.println("current time: " + currentTime
					// + " time interval: " + timeInterval);
					if (currentTime + timeInterval >= totalTime) {
						currentTime = totalTime;
						// System.out.println("current time is: " +
						// currentTime);
					} else {
						// System.out.println("current time is: " +
						// currentTime);
						currentTime = currentTime + timeInterval;
						// System.out.println("after addition, current time is: "
						// + currentTime);
						currentState = sampleNextNucleotideState(currentState);
						sampledPath.add(new Integer(currentState));
						sampledTime.add(new Double(currentTime));
					}

				}
				endState = sampledPath.get(sampledPath.size() - 1).intValue();
			}
			sampledPath.add(new Integer(childState));
			sampledTime.add(new Double(totalTime));
		} else {
			while (endState != childState) {
				// initialize path
				sampledPath.clear();
				sampledPath.add(new Integer(parentState));
				sampledTime.clear();
				sampledTime.add(new Double(0.0));

				// initialize time and state
				currentTime = 0.0;
				currentState = parentState;
				boolean firstSample = true;
				while (currentTime < totalTime) {
					if (firstSample) {
						timeInterval = sampleFirstSubstitutionTime(totalTime);
						// System.out.println("current time: " + currentTime
						// + " time interval: " + timeInterval);
						currentTime = currentTime + timeInterval;
						currentState = sampleNextNucleotideState(currentState);
						sampledPath.add(new Integer(currentState));
						sampledTime.add(new Double(currentTime));
						firstSample = false;
					} else {
						timeInterval = randGenerator.nextExponential(1.0);
						// System.out.println("current time: " + currentTime
						// + " time interval: " + timeInterval);
						if (currentTime + timeInterval >= totalTime) {
							currentTime = totalTime;
						} else {
							currentTime = currentTime + timeInterval;
							currentState = sampleNextNucleotideState(currentState);
							sampledPath.add(new Integer(currentState));
							sampledTime.add(new Double(currentTime));
						}
					}
				}
				endState = sampledPath.get(sampledPath.size() - 1).intValue();
			}
			sampledPath.add(new Integer(childState));
			sampledTime.add(new Double(totalTime));
		}

		return (sampledPath.size() - 2);
		// System.out.println(sampledPath.toString() + " size: "
		// + sampledPath.size());
		// System.out.print((sampledPath.size() - 2) + ";");
		// System.out.print((sampledPath.size() - 2) + ";");
	}

	// sample next nucleotide state based on current state
	// JC69 used
	// if use other substitution model, this function needs to be changed
	public int sampleNextNucleotideState(int currentState) {
		// create an integer to store final result
		int sampledNucleo = -1;

		List<Integer> possibleStates = new ArrayList<Integer>();
		int[] allStates = new int[] { 0, 1, 2, 3 };
		for (int i : allStates) {
			if (i != currentState) {
				possibleStates.add(new Integer(i));
			}
		}

		// draw a random number between 0 and 1
		double randNum = generator.nextDouble();

		// determine which nucleotide to return
		if (randNum < (1.0 / 3.0)) {
			sampledNucleo = possibleStates.get(0);
		} else if ((randNum >= (1.0 / 3.0)) && (randNum < (2.0 / 3.0))) {
			sampledNucleo = possibleStates.get(1);
		} else {
			sampledNucleo = possibleStates.get(2);
		}
		return sampledNucleo;
	}

	public double sampleFirstSubstitutionTime(double totalTime) {
		// store the final result
		double sampledFirstTime;

		// draw a random number between 0 and 1
		double randNum = generator.nextDouble();

		// calculate the time for the first substitution
		sampledFirstTime = -Math
				.log(1.0 - randNum * (1 - Math.exp(-totalTime)));
		return sampledFirstTime;
	}

}